from __future__ import print_function
import numpy as np
import scipy.signal as signal
import matplotlib.pyplot as plt


def vad(x, frame_size=None, sr=None, frame_shift=None, plot_seg=False):
    """
    计算vad值
        :param x: 一段时间段
        :param frame_size:
        :param sr: 目标取样率
        :param frame_shift:
        :param plot_seg: 是否画图
        :return:
    """
    if sr is None:
        sr = 16000
    if frame_size is None:
        frame_size = 256
    if frame_shift is None:
        frame_shift = 128
    amp_th1 = 8
    amp_th2 = 20
    zcr_th = 5

    max_silence = 8
    min_len = 15
    status = 0
    count = 0
    silence = 0

    x = x/np.absolute(x).max()

    tmp1 = enframe(x[0:(len(x)-1)], frame_size, frame_shift)
    tmp2 = enframe(x[1:(len(x)-1)], frame_size, frame_shift)
    signs = (tmp1 * tmp2) < 0
    diffs = (tmp1 - tmp2) > 0.05
    zcr = np.sum(signs * diffs, axis=1)

    filter_coeff = np.array([1, -0.9375])
    pre_emphasis = signal.convolve(x, filter_coeff)[0:len(x)]
    amp = np.sum(np.absolute(enframe(pre_emphasis, frame_size, frame_shift)), axis=1)

    amp_th1 = min(amp_th1, amp.max() / 3)
    amp_th2 = min(amp_th2, amp.max() / 8)

    x1 = []
    x2 = []
    t = 0

    for n in range(len(zcr)):
        if status == 0 or status == 1:
            if amp[n] > amp_th1:
                x1.append(max(n - count - 1, 1))
                status = 2
                silence = 0
                count = count + 1
            elif amp[n] > amp_th2 or zcr[n] > zcr_th:
                status = 1
                count = count + 1
            else:
                status = 0
                count = 0
            continue
        if status == 2:
            if amp[n] > amp_th2 or zcr[n] > zcr_th:
                count = count + 1
            else:
                silence = silence + 1
                if silence < max_silence:
                    count = count + 1
                elif count < min_len:
                    status = 0
                    silence = 0
                    count = 0
                else:
                    status = 0
                    count = count - silence / 2
                    x2.append(x1[t] + count - 1)
                    t = t + 1

    # 画图
    if plot_seg:
        plt.figure('speech endpoint detect')
        plt.plot(np.arange(0, len(x)) / (float)(sr),x, "b-")
        len_endpoint = min(len(x1), len(x2))
        for i in range(len_endpoint):
            plt.vlines(x1[i] * frame_shift / (float)(sr), -1, 1, colors="c", linestyles="dashed")
            plt.vlines(x2[i] * frame_shift / (float)(sr), -1, 1, colors="r", linestyles="dashed")
        plt.xlabel("Time/s")
        plt.ylabel("Normalized Amp")
        plt.grid(True)
        plt.show()
    return x1, x2


def enframe(x, frame_size, frame_shift):
    """

        :param x:
        :param frame_size:
        :param frame_shift:
        :return:
    """
    x_len = len(x)
    nf = (int)((x_len - frame_size + frame_shift) / frame_shift)
    # f = np.zeros((nf, frame_size), dtype=np.float32)
    ind_f = frame_shift * (np.arange(0, nf)).reshape(nf, 1)
    ind_s = np.arange(0, frame_size).reshape(1, frame_size)
    ind_all = np.tile(ind_f, (1, frame_size)) + np.tile(ind_s, (nf, 1))
    f = x[ind_all]
    return f

